package com.coocaa.typhoon.api.v1.milvus.service;

import cn.hutool.core.bean.BeanUtil;
import com.coocaa.cirrus.common.base.util.BeanConvertUtil;
import com.coocaa.cirrus.common.base.util.DateUtil;
import com.coocaa.typhoon.api.third.volcengine.constant.EmbeddingTypeConstant;
import com.coocaa.typhoon.api.third.volcengine.model.EmbeddingQueryDTO;
import com.coocaa.typhoon.api.third.volcengine.provider.VolcEmbeddingProvider;
import com.coocaa.typhoon.api.third.volcengine.service.TOSService;
import com.coocaa.typhoon.api.v1.milvus.constant.MilvusConstant;
import com.coocaa.typhoon.api.v1.milvus.convert.MilvusFilterConverter;
import com.coocaa.typhoon.api.v1.milvus.model.HybridSearchCondition;
import com.coocaa.typhoon.api.v1.milvus.model.UserFileMilvusDTO;
import com.coocaa.typhoon.api.v1.milvus.model.UserFileMilvusVO;
import com.coocaa.typhoon.api.v1.userfile.util.TosFinalUrlUtil;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.SortedMap;

/**
 * @author bijiahao
 * @date : 2025/9/25 10:48
 * @description
 */
@Slf4j
@Service
public class UserFileMilvusQueryService {
    @Autowired
    private MilvusService milvusService;
    @Autowired
    private VolcEmbeddingProvider volcEmbeddingProvider;
    @Autowired
    private TOSService tosService;

    private HybridSearchCondition buildCondition(String userId, String text, List<String> fileTypeList, String location, Integer top, List<String> tosLinkList, String startTime, String endTime) {
        SortedMap<Long, Float> questionSparseVector = null;
        EmbeddingQueryDTO embeddingQueryDTO = new EmbeddingQueryDTO();
        embeddingQueryDTO.setContent(text);
        embeddingQueryDTO.setType(EmbeddingTypeConstant.TEXT);
        List<Float> questionVector = volcEmbeddingProvider.questionVector(embeddingQueryDTO);
        if (null == questionSparseVector && null == questionVector) {
            return null;
        }
        List<String> cols = BeanConvertUtil.getColumnNames(UserFileMilvusDTO.class);
        int denseTopK = top;
        int sparseTopK = top;
        String expr = MilvusFilterConverter.getFilter(userId, fileTypeList, DateUtil.toTimestamp(startTime), location, DateUtil.toTimestamp(endTime));
        HybridSearchCondition condition = HybridSearchCondition.builder()
                .collectionName(MilvusConstant.USER_FILE)
                .queryDense(questionVector)
                .querySparse(questionSparseVector)
                .vectorFieldName(MilvusConstant.VECTOR_FIELD_NAME)
                .denseTopK(denseTopK)
//                .denseParams("{\"nprobe\": 10}")
//                .denseParams("{\"nprobe\": 32}")
                .denseParams("{\"ef\": 64}")
                .sparseVectorFieldName(MilvusConstant.SPARSE_VECTOR_FIELD_NAME)
                .sparseTopK(sparseTopK)
                .sparseParams("{\"drop_ratio_search\": 0.2}")
                .rerankerList(Arrays.asList(0.7f, 0.3f))
                .topK(top)
                .outFields(cols)
                .expr(expr)
                .build();
        return condition;
    }

    public List<UserFileMilvusVO> hybridSearch(String userId,String text, List<String> fileTypeList, String location, Integer top, List<String> tosLinkList, String startTime, String endTime) {
        long start = System.currentTimeMillis();
        HybridSearchCondition condition = buildCondition(userId, text, fileTypeList, location, top, tosLinkList, startTime, endTime);
        long end1 = System.currentTimeMillis();
        log.info("buildCondition cost: {}秒", (end1 - start) / 1000.0f);
        if (null == condition) {
            return new ArrayList<>();
        }
        List<List<SearchResp.SearchResult>> searchResults = milvusService.hybridSearch(condition);
        List<UserFileMilvusVO> res = new ArrayList<>();
        searchResults.forEach(searchResult -> {
            searchResult.stream().forEach(row -> {
                Float score = row.getScore();
                UserFileMilvusDTO milvusDTO = BeanConvertUtil.convertMapToEntity(row.getEntity(), UserFileMilvusDTO.class);
                UserFileMilvusVO tmpRes = new UserFileMilvusVO();
                BeanUtil.copyProperties(milvusDTO, tmpRes);
                tmpRes.setScore(score);
                tmpRes.setFileUrl(tosService.getSignTosHttpUrl(tmpRes.getFilePath()));
                res.add(tmpRes);
            });
        });
        long end = System.currentTimeMillis();
        log.info("hybridSearch cost: {}秒", (end - end1) / 1000.0f);
        return res;
    }
}
